package com.trie;

import org.springframework.util.StringUtils;

import java.util.HashMap;
import java.util.Objects;

/**
 * @Author: tgy
 * @Date: 2020-10-12 13:55
 */
public class HashTrie<V> implements Trie<V> {

    private int size;

    private Node<V> root;

    private static class Node<V> {

        HashMap<Character,Node<V>> childNode;
        V value;
        boolean isEnd;

        Node<V> parent;

        public Node(Node<V> node) {

            parent = node;
            childNode = new HashMap<>();
            isEnd = false;
        }
    }


    public HashTrie() {

        root = new Node<>(null);
        size = 0;
    }

    @Override
    public int size() {
        return size;
    }

    @Override
    public boolean isEmpty() {
        return size == 0;
    }

    @Override
    public void clear() {

        root = null;
        size = 0;
    }

    @Override
    public boolean contains(String key) {

        if (StringUtils.isEmpty(key)) {

            return false;
        }

        Node<V> node = nodeOfKey(key);
        return Objects.nonNull(node) && node.isEnd;
    }

    @Override
    public V add(String key, V value) {

        checkKey(key);

        Node<V> node = root;
        Node<V> newNode;
        Node<V> childNode;
        boolean isNodeNull = false;
        for (int i = 0; i < key.length(); i++) {

            char charKey = key.charAt(i);

            if (!isNodeNull) {

                childNode = node.childNode.get(charKey);
                if (Objects.isNull(childNode)) {
                    isNodeNull = true;
                }else {
                    node = childNode;
                }
            }

            if (isNodeNull) {

                newNode = new Node<>(node);
                node.childNode.put(charKey,newNode);
                node = newNode;
            }
        }

        node.isEnd = true;
        V oldVal = node.value;
        node.value = value;

        size++;
        return oldVal;
    }

    @Override
    public V get(String key) {

        Node<V> node = nodeOfKey(key);

        return Objects.isNull(node)?null:node.value;
    }

    @Override
    public V remove(String key) {

        Node<V> node = nodeOfKey(key);

//        不存在的key, 直接返回
        if (Objects.isNull(node) || !node.isEnd) {

            return null;
        }

        size--;
        V removeNode = null;
//        如果还有子节点
        if (!node.childNode.isEmpty()) {

            node.isEnd = false;
            removeNode = node.value;
            node.value = null;
            return removeNode;
        }

        removeNode = node.value;
        node.isEnd = false;

        for (int i = key.length() - 1; i >= 0; i--) {

            Node<V> pNode = node.parent;

            if (!node.isEnd && node.childNode.isEmpty()) {

                pNode.childNode.remove(key.charAt(i));
            }else {

                break;
            }

            node = pNode;
        }

        return removeNode;
    }

    @Override
    public boolean startWith(String prefix) {

       checkKey(prefix);

        return Objects.nonNull(nodeOfKey(prefix));
    }


    /**
     * 通过 key 获取 Node<V>
     * @param key by key
     * @return
     */
    private Node<V> nodeOfKey(String key) {

        checkKey(key);

        Node<V> node = root;

        for (int i = 0; i < key.length(); i++) {

            Node<V> childNode = node.childNode.get(key.charAt(i));

            if (Objects.isNull(childNode)) {

                return null;
            }

            node = childNode;
        }

        return node;
    }


    /**
     * 检查key 是否非法
     * @param key
     */
    private void checkKey(String key) {

        if (StringUtils.isEmpty(key)) {

            throw new IllegalArgumentException("key not null or empty");
        }
    }
}
